from detect import *

import argparse

from models.yolo import Model
from datasets import *
from utils.utils import *
import ctypes

libLoad = ctypes.cdll.LoadLibrary
sharelib = libLoad("/media/st/Application/Ubuntu/NRSL/object_detection/history/img_depth_ws/devel/lib/libpython_share_memory.so") #need change 

def get_img():
    sharelib.get_img.restype =ctypes.POINTER(ctypes.c_uint8)  #c_char_p
    pointer_img = sharelib.get_img()
    size=(480,640,3)
    RecImg = np.ctypeslib.as_array(pointer_img,shape=size)
    # RecImg =  np.reshape(np.array(np.fromiter(pointer_img, dtype=np.uint8, count=size[0]*size[1]*size[2])) ,size)
    return RecImg
 
def get_img_flag():
    sharelib.get_img_flag.restype = ctypes.c_int
    return sharelib.get_img_flag()

def send_img_pre_data(mea_data):
    if mea_data is not None:
        pre_class =  np.asarray(mea_data[:, 5] , dtype=np.int32)
        pre_2dbbox =  np.asarray(mea_data[:, 0:4], dtype=np.int32)
        if not pre_class.flags['C_CONTIGUOUS']:
            pre_class = np.ascontiguous(pre_class, dtype=pre_class.dtype)  
        p_pre_class = ctypes.cast(pre_class.ctypes.data, ctypes.POINTER(ctypes.c_int))   

        if not pre_2dbbox.flags['C_CONTIGUOUS']:
            pre_2dbbox = np.ascontiguous(pre_2dbbox, dtype=pre_2dbbox.dtype)
        p_pre_2dbbox = ctypes.cast(pre_2dbbox.ctypes.data, ctypes.POINTER(ctypes.c_int))   
        sharelib.send_img_pre_data(p_pre_class, p_pre_2dbbox, pre_class.shape[0])
    else:
        pre_class =  np.zeros((1,1), dtype=np.int32)
        pre_2dbbox =  np.zeros((1,4), dtype=np.int32)
        if not pre_class.flags['C_CONTIGUOUS']:
            pre_class = np.ascontiguous(pre_class, dtype=pre_class.dtype)  
        p_pre_class = ctypes.cast(pre_class.ctypes.data, ctypes.POINTER(ctypes.c_int))   

        if not pre_2dbbox.flags['C_CONTIGUOUS']:
            pre_2dbbox = np.ascontiguous(pre_2dbbox, dtype=pre_2dbbox.dtype)
        p_pre_2dbbox = ctypes.cast(pre_2dbbox.ctypes.data, ctypes.POINTER(ctypes.c_int))   
        sharelib.send_img_pre_data(p_pre_class, p_pre_2dbbox, 0)


class yolov5:
    def __init__(self, model_file, model_cfg, imagesz = 640):
        self.weights = model_file
        self.imgsz = imagesz
        self.device = torch_utils.select_device('0') 
        self.augment = False
        self.conf_thres = 0.7
        self.iou_thres = 0.5
        self.classes = None
        self.agnostic_nms = False
        self.model = Model(model_cfg).to(self.device)
        # self.model = torch.load(self.weights, map_location=self.device)['model']
        self.model.load_state_dict(torch.load(self.weights))
        self.model.to(self.device).eval()
        img = torch.zeros((1, 3, self.imgsz, self.imgsz), device=self.device)  # init img
        _ = self.model(img.float()) if self.device.type != 'cpu' else None  # run once

    def detect(self, image_raw):
        image_det = image_raw.copy()
        image_det = letterbox(image_det, new_shape=self.imgsz)[0]
        image_det = image_det[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        image_det = np.ascontiguousarray(image_det)

        image_det = torch.from_numpy(image_det).to(self.device)
        image_det = image_det.float()  # uint8 to fp16/32
        image_det /= 255.0  # 0 - 255 to 0.0 - 1.0
        if image_det.ndimension() == 3:
            image_det = image_det.unsqueeze(0)
        
        # Inference
        pred = self.model(image_det, augment=self.augment)[0]

        # Apply NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres,\
            fast=True, classes=self.classes, agnostic=self.agnostic_nms)
        det = pred[0]
        # gn = torch.tensor(image_raw.shape)[[1, 0, 1, 0]]  #  normalization gain whwh
        if det is not None and len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_coords(image_det.shape[2:], det[:, :4], image_raw.shape).round()  # x1y1x2y2,  conf, cls
            
        return det

class Image_detection: 
    def __init__(self ):
        self.image_detector = yolov5("/media/st/Application/Ubuntu/NRSL/object_detection/history/yolov5_learning/weights/best_weight.pt", \
            "/media/st/Application/Ubuntu/NRSL/object_detection/history/yolov5_learning/models/score/fire_danger.yaml", 640)

    def img_detection(self, cv_image):
        det = self.image_detector.detect(cv_image)
        # det = self.image_detector.detect(cv_image)
        if det is None or len(det) == 0:
            return

        det = det.detach().cpu().numpy()
        # print("det type ", det)
        return det

if __name__ == "__main__":    
    image_det=Image_detection()
    print("yolov5 launch!")
    while True:
        # time.sleep(0.1)
        # print("img flag ", get_img_flag())
        if get_img_flag()==1 :
            t0 = time.time()
            img = get_img()
            send_img_pre_data(image_det.img_detection(img))
            torch.cuda.empty_cache()
            t1= time.time() - t0
            print("time ", t1)

